#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""

"""

from pprint import pprint
import copy

from linear_tools import linear_combination
from matrix_tools import multiply, multiply_1, identity, matrix_append_col
# from linear_utils import eliminate_vector


class NoSuitableLine(Exception): pass

class Gaussian:

    # =================== 两个初等变换 =====================
    def _row_switch_trans(self, r1, r2, mat=None):
        ''' 交换行 '''
        if not mat:
            mat = self.matrix

        mat[r1], mat[r2] = mat[r2], mat[r1]

    def _eliminate_trans(self, row_idx, pivot_idx, matrix=None):
        ''' k * ri  + 1 * rj '''
        if not matrix:
            matrix = self.matrix

        pivot_row = matrix[pivot_idx]
        cur_row   = matrix[row_idx]

        coefficient = cur_row[pivot_idx] / pivot_row[pivot_idx]
        matrix[row_idx] = list( linear_combination(
            [pivot_row, cur_row],
            [-coefficient, 1]
        ))
        return coefficient

    def _not_zero(self, r, c):
        return 0 != self.matrix[r][c]

    def _is_zero(self, r, c):
        return 0 == self.matrix[r][c]

    def _record_switch_transform(self, r1, r2):
        print('=================== switch =====================')
        self.transform_list.append('E(%s <--> %s)' % (r1, r2))

        # self._row_switch_trans(row_idx, cand_row, self.elementary_trans)
        self._row_switch_trans(r1, r2, self.elementary_trans)
        self._row_switch_trans(r1, r2, self.elementary_switch)

    def _record_elimit_transform(self, pivot_row, cur_row, co):
        '''
            pivot_row: 元来自第几行
            cur_row: 根据pivot_row 消第几行
        '''
        self.transform_list.append('M: (%s * %s) + (1 * %s) ' % (-co, pivot_row, cur_row))

        # 记录初等变换矩阵
        ele_trans = self.elementary_trans
        ele_trans[cur_row] = list( linear_combination(
            [ele_trans[pivot_row], ele_trans[cur_row]],
            [-co, 1]
        ))

    def _pivoting(self, row_idx):
        ''' 如果当前行pivot位置为0, 找到当前行之后的行中pivot_idx/row_idx 不为0 的行, 并和当前行交换.
            如果没有找到, 说明处理的矩阵为singular矩阵 raise Exception.
            这个寻找合适的pivot行的过程叫做pivoting
        '''
        pivot_col_idx = row_idx
        if self._not_zero(row_idx, pivot_col_idx):
            return

        for cand_row in range(row_idx+1, self.rows):
            if self._not_zero(cand_row, pivot_col_idx):
                self._row_switch_trans(row_idx, cand_row)
                self._record_switch_transform(row_idx, cand_row)
                return

        raise NoSuitableLine()

    def _eliminate_line(self, process_row_idx, pivot_idx):
        if self._is_zero(process_row_idx, pivot_idx):
            return

        coefficient = self._eliminate_trans(process_row_idx, pivot_idx)

        # 记录消元的初等变换过程
        self._record_elimit_transform(pivot_idx, process_row_idx, coefficient)

    def _normalization(self, row_idx):
        ''' 对 (row_idx, pivot_idx==row_idx) 做归一化处理

            示例输入: 2 2 3
                      0 2 5
                      0 0 6

            示例输出: 2 2 3 ------> 1, 1, 1.5
                      0 2 5 ------> 0, 1, 2.5
                      ...
        '''
        pivot_idx = row_idx
        co = self.matrix[row_idx][pivot_idx]
        assert co
        if co == 1:
            return

        self.matrix[row_idx] = [value/co for value in self.matrix[row_idx]]

    def _eliminate_pivot(self, pivot_idx):
        ''' 处理第几个pivot '''
        self._pivoting(pivot_idx)    # 保证pivot_idx维护行这个位置不为0
        self._normalization(pivot_idx)

        # for row_idx in range(pivot_idx+1, self.rows):
        for row_idx in range(0, self.rows):
            '''  NOTE: 这里大胆的把当前行之前的行和当前行之后的行都消去, 省掉了最后一步back-substitution '''
            if row_idx == pivot_idx: # 忽略当前行
                continue
            self._eliminate_line(row_idx, pivot_idx)

    # def _back_sub(self):
    #     pass

    def _uable_pivot(self, pivot_idx):
        self._pivoting(pivot_idx)    # 保证pivot_idx维护行这个位置不为0

        for row_idx in range(pivot_idx+1, self.rows):
            self._eliminate_line(row_idx, pivot_idx)

    def _init_matrix(self, matrix, inplace=True):
        ''' 对aug_matrix 增光矩阵进行高斯消元, 可以用来求逆矩阵, 也可以解方程组 '''
        if inplace:
            self.matrix = matrix
        else:
            self.matrix = copy.deepcopy(matrix)
        self.rows = len(matrix)
        self.cols = len(matrix[0])
        assert self.cols >= self.rows, (self.rows, self.cols)

        self.transform_list = []                     # 语言描述消元过程
        self.elementary_switch = identity(self.rows) # 初等交换矩阵: 每次初等变换, 记录交换两行的操作, 记为S
        self.elementary_trans = identity(self.rows)  # 初等变换矩阵, 随着每次初等变换操作, 最后变为EA = U 的"E"(不包括交换操作)

    def show(self, mat=None):
        if not mat:
            mat = self.matrix

        print('=================== show matrix =====================')
        show(mat)

    # =================== outside interface =====================
    def rref(self, mat):
        ''' 获得矩阵的reduced echelong form '''
        # self.

    def U(self, matrix):
        ''' 把矩阵变为upper triangle 矩阵 '''
        self._init_matrix(matrix)
        try:
            list( map(self._uable_pivot, range(self.rows)) )
        except Exception as e:
            print('catched exception %s' % type(e))
            raise e

        pprint(self.transform_list)
        return self.elementary_switch, self.elementary_trans, self.matrix

    def gaussian_eliminate(self, aug_matrix):
        ''' 对aug_matrix 增光矩阵进行高斯消元, 可以用来求逆矩阵, 也可以解方程组 '''
        self._init_matrix(aug_matrix)
        try:
            # list只是为了让 generator 执行, 不然只是一个生成器(没有真正执行)
            list( map(self._eliminate_pivot, range(self.rows)) )
        except Exception as e:
            print('catched exception %s' % type(e))
            raise e

        return self.transform_list

    def I(self, mat):
        ''' 求矩阵的逆 Read More: numpy.mat.I '''
        rows = len(mat)
        cols = len(mat[0])
        assert rows == cols

        iden = identity(rows)
        new_mat = matrix_append_col(mat, iden)
        self.gaussian_eliminate(new_mat)

        inverse_mat = []
        for r in new_mat:
            inverse_mat.append(r[-rows:])

        return inverse_mat

    def four_fourmental_basis(self, mat):
        '''
            求解矩阵的四个空间的基.
            column space
            null space

            row space
            left space
        '''

def show(mat, msg=''):
    print('=================== mat: %s =====================' % msg)
    for r in mat:
        print(r)
    # print('=================== show end =====================')

def test_inverse():
    mat = [
        [1, 2, 8, 9],
        [3, 4, 3, 7],
        [9, 8, 5, 13],
        [8, 9, 7, 13],
    ]
    # =================== numpy mat inverse  =====================
    import numpy as np
    n = 10
    # mat = np.random.rand(n, n)
    np_mat = np.mat(mat)
    np_inverse = np.mat(np_mat).I
    show(np_inverse)

    # =================== my mat inverse =====================
    g = Gaussian()
    my_inverse = g.I(mat)
    show(my_inverse, 'my inverse')

def test_matrix():
    print('test_matrix')
    g = Gaussian()
    # A = [
    #     [1, 2, 8, 4],
    #     [3, 4, 3, 9],
    #     [9, 8, 5, 6],
    # ]
    A = [
        [1, 2, 8, 9],
        [10, 0, 3, 7],
        [9, 8, 5, 13],
        [8, 9, 7, 13],
    ]
    # U = copy.deepcopy(A)

    S, E, U = g.U(A)               # switch, elementary
    show(A, "A")           # A
    show(U, "U")           # U
    # show(S, "S")           # S
    show(E, "E")           # E

    EA = multiply_1(E, A)  # EA
    show(EA, "E * A == U") # E * S * A = U
    # ESA = multiply(E, S, A)  # ESA
    # show(ESA, "E * S * A == U") # E * S * A = U

    # L 具有很好的性质:transform_list 和L相对应.
    # transform_list反而和E对应的不是那么美观.
    # TODO1: 根据transform_list求解L
    L = g.I(E)              # L
    show(L, "L")                    # L

    LU = multiply_1(L, U)           # LU
    show(LU, "L * U == A")          # LU = A

    # A = LU ===> AI=UI*LI      # U逆 * L逆得到A逆
    # TODO: UI 和 LI 更快的算法, 因为都是三角矩阵,所以应该有的.
    UI = g.I(U)
    LI = g.I(L)
    UI_LI = multiply_1(UI, LI)
    show(UI_LI, "UI * LI == AI")          # UI *LI = AI

    AI = g.I(A)                   # A的逆
    show(AI, "AI, A's inverse")     # LU = A

    U_AI = multiply_1(U, AI)        # U_AI
    show(U_AI, "U * AI == E")       # LU = A

    # =================== inverse to unknows  =====================
    # A * unknows = values
    # (AI * A) * unknows = AI * values
    values = [
        [15, 22],
        [33, 42],
        [99, 88],
        [87, 92],
    ]

    unknows = multiply_1(AI, values)
    show(unknows, 'unknows')

    # g eliminate values
    new_mat = matrix_append_col(A, values)
    g.gaussian_eliminate(new_mat)
    show(new_mat, 'eliminate to solve unknows')

def test_change_basis():
    mat = [
        [1, 3],
        [2, 4],
    ]
    g = Gaussian()
    i = g.I(mat)
    show(i, 'basis inverse')
    t = multiply([[2, 6]], i)
    show(t, 't')
    t = multiply([[4, 8]], i)
    show(t, 't')
    t = multiply([[3, 7]], i)
    show(t, 't')

def test():
    # test_inverse()
    test_matrix()
    # test_change_basis()

def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
